"""会议安排奖励计算模块"""


class MeetingReward:
    """会议安排任务的奖励计算类"""
    
    def __init__(self, **kwargs):
        """初始化奖励计算器"""
        self.reward = None
    
    def format_reward(self, response: str) -> float:
        """
        格式奖励：检查响应格式是否正确
        
        Args:
            response: LLM的响应文本
            
        Returns:
            格式奖励分数
        """
        # 检查是否包含必要的标签
        reward = 0
        if '<think>' in response and '</think>' in response:
            reward += 0.5
        if '<action>' in response and '</action>' in response:
            reward += 0.5
        if '<invoke>' in response and '</invoke>' in response:
            reward += 0.5
        
        return reward
    
    def response_reward(self, response: str, reward_client=None, resp_sys_prompt: str = None) -> float:
        """
        响应质量奖励（可选：使用奖励模型评估）
        
        Args:
            response: LLM的响应
            reward_client: 奖励模型客户端
            resp_sys_prompt: 奖励模型的系统提示词
            
        Returns:
            响应质量奖励分数
        """
        # 这里可以实现基于规则或模型的响应质量评估
        return 0
    
    def action_reward(self, action: dict, state: dict) -> float:
        """
        动作奖励：评估动作的合理性
        
        Args:
            action: 执行的动作
            state: 当前状态
            
        Returns:
            动作奖励分数
        """
        tool_name = action.get('tool_name', '')
        
        # 鼓励按照合理流程执行
        rewards_map = {
            'query_calendar': 1.0,      # 鼓励先查询日历
            'query_history': 1.0,       # 鼓励查询历史
            'negotiate_time': -2.0,     # 协商有成本
            'select_time': 0,           # 选择时间不额外奖励
            'calculate': 0.5            # 计算工具
        }
        
        return rewards_map.get(tool_name, 0)
    
    def final_reward(self, state: dict, selected_time: dict) -> float:
        """
        最终奖励：评估会议安排的综合质量
        
        Args:
            state: 最终状态
            selected_time: 选择的会议时间
            
        Returns:
            最终奖励分数
        """
        reward = 0
        
        if not selected_time:
            return -100  # 未选择时间，重罚
        
        day = selected_time.get('day', '')
        time_slot = selected_time.get('time_slot', '')
        
        # 从状态中提取信息
        required_attendees = state.get('required_attendees', [])
        optional_attendees = state.get('optional_attendees', [])
        availability = state.get('availability', {})
        historical_pattern = state.get('historical_pattern', {})
        negotiation_history = state.get('negotiation_history', [])
        
        # 1. 计算必选人员参与率
        required_present = 0
        for person in required_attendees:
            if person in availability:
                if day in availability[person]:
                    if availability[person][day].get(time_slot, 'busy') == 'available':
                        required_present += 1
        
        required_rate = required_present / max(1, len(required_attendees))
        if required_rate < 1.0:
            missing = len(required_attendees) - required_present
            reward -= missing * 100  # 每缺少一个必选人员扣100分
        
        # 2. 计算可选人员参与率
        optional_present = 0
        for person in optional_attendees:
            if person in availability:
                if day in availability[person]:
                    if availability[person][day].get(time_slot, 'busy') == 'available':
                        optional_present += 1
        
        optional_rate = optional_present / max(1, len(optional_attendees))
        reward += optional_present * 1  # 每个可选人员参加加1分
        
        # 3. 协商成本
        negotiations_cnt = len(negotiation_history)
        reward -= negotiations_cnt * 2  # 每次协商扣2分
        
        # 4. 历史模式匹配奖励
        if historical_pattern:
            most_common_day = historical_pattern.get('most_common_day', '')
            most_common_slot = historical_pattern.get('most_common_slot', '')
            
            if day == most_common_day and time_slot == most_common_slot:
                frequency = historical_pattern.get('frequency', 0) / 100.0
                reward += frequency * 5  # 历史匹配最高加5分
        
        return reward
    
    def step_reward(self, action: dict, result: str, state: dict) -> float:
        """
        步骤奖励：每一步动作的即时奖励
        
        Args:
            action: 执行的动作
            result: 动作执行结果
            state: 当前状态
            
        Returns:
            步骤奖励分数
        """
        tool_name = action.get('tool_name', '')
        
        # 查询动作有小额正奖励
        if tool_name in ['query_calendar', 'query_history']:
            return 0.5
        
        # 协商动作有成本
        elif tool_name == 'negotiate_time':
            if 'success' in result.lower():
                return -2  # 成功的协商仍有成本
            else:
                return -5  # 失败的协商成本更高
        
        # 计算动作
        elif tool_name == 'calculate':
            return 0
        
        return 0

